#!/usr/bin/env python

###############################################################################
#
# Use masks to count storm tracks within ocean basins.
#
###############################################################################

#
# MODULES
#
import sys
import os
import re
import glob
import string

#from numarray import *
#from pylab import *
#import Numeric

#import Nio #  for reading netCDF files
import cdms, cdutil

#
# ARGUMENTS
#

pf = 'past'

if len(sys.argv) < 2:
    print 'Usage: %s <past|future>' % sys.argv[0]
    raise SystemExit
else:
   pf = sys.argv[1]

# input mask file
maskfilename = '/work/romano/fvcam2.2/oceanbasin_masks.nc'
#maskfile = Nio.open_file(maskfilename, "r")
maskfile = cdms.open(maskfilename)

# dictionary of (basin, count) and (basin, mask)
tccounts = dict()
masks = dict()

runs = [2, 3, 4, 5, 6, 7, 8, 10, 11, 13, 14, 15, 17, 18]
basins = ['north', 'northatlantic', 'northindian', 'northpacific',
          'south', 'southindian', 'southpacific']

for basin in basins:
    masks[basin] = maskfile(basin)
# close maskfile

bounds_lon = maskfile('bounds_lon')
bounds_lat = maskfile('bounds_lat')

start_lon = float(bounds_lon[0][0])
start_lat = float(bounds_lat[0][0])

step_lon = bounds_lon[1][1]-bounds_lon[1][0]
step_lat = bounds_lat[1][1]-bounds_lat[1][0]


nlon = bounds_lon.shape[0]
nlat = bounds_lat.shape[0]

# loop through runs
for nrun in runs:
    print nrun
    for basin in basins:
        tccounts[basin] = 0

    fname = '%d/%s/all/traj_%d_%s_summary' % (nrun, pf, nrun, pf)

    if os.path.exists(fname):
        for line in open(fname).readlines():

            # get latitude and longitude
            elems = string.split(line)

            if not elems[0].startswith('ID'): # ignore header
                lon = float(elems[5])
                lat = float(elems[6])

                # guess index and then search
                lon_start = int(round((lon-start_lon)/step_lon)-1)
                lat_start = int(round((lat-start_lat)/step_lat)-1)

                # find lon/lat indices from bounds arrays
                i = lon_start
                cont=0
                while i < nlon and cont==0:
                    if lon>=float(bounds_lon[i][0]) and lon<float(bounds_lon[i][1]):
                        j = lat_start
                        while j < nlat and cont==0:
                            if lat>=float(bounds_lat[j][0]) and lat<float(bounds_lat[j][1]):
                                ilon = i
                                ilat = j
                                # loop through basins and increment each count
                                for basin in basins:
                                    basinmask = masks[basin]
                                    if float(basinmask[ilat][ilon]) == 1.0:
                                        tccounts[basin] = tccounts[basin] + 1
                                cont=1
                            j=j+1
                    i=i+1
        
    # write out counts
    outfilename = 'basin/tccount_%d_%s.txt' % (nrun, pf)
    print outfilename
    outfile = open(outfilename,'w')
    for basin in tccounts:
        bcount = '%s %d\n' % (basin, tccounts[basin])
        outfile.write(bcount)
    outfile.close()    


